package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
	"ntrace/enviroments"
	"ntrace/utils/routines"
	"time"
)

type TCPSYNReceiver struct {
	BaseReceiver
	TCPConn *ipv4.RawConn
}

func NewTCPSYNReceiver(ctx context.Context, task *TraceTask, stat *TraceStat) *TCPSYNReceiver {
	return &TCPSYNReceiver{
		BaseReceiver: BaseReceiver{
			Ctx:  ctx,
			Name: "ICMPReceiver",
			Task: task,
			Stat: stat,
		},
	}
}

func (r *TCPSYNReceiver) StartReceiver() error {
	routines.GoWithRecover("TCPSYNReceiver-icmp", r.doReceiveICMP)
	routines.GoWithRecover("TCPSYNReceiver-tcp", r.doReceiveTCP)
	return nil
}

func (r *TCPSYNReceiver) StopReceiver() {
	r.BaseReceiver.StopReceiver()
	if r.TCPConn != nil {
		r.TCPConn.Close()
	}
}

func (r *TCPSYNReceiver) doReceiveICMP() {
	golog.Debugf("UDPReceiver doReceiveICMP")
	conn, err := OpenReceiverIP4PacketConn("ip4:icmp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	filter := CreateBPFTCPICMP(net.ParseIP(r.Task.DstIP))
	err = SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():

			conn.Close()
			return
		default:
			r.receiveICMPPkt(payloadBuf)
		}

	}
}

func (r *TCPSYNReceiver) receiveICMPPkt(payloadBuf []byte) {
	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(time.Second))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("read packet failed. err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	icmpMsg, err := icmp.ParseMessage(1, payload)

	devIP := h.Src.To4().String()
	var originIPPkt []byte
	switch icmpMsg.Type {
	case ipv4.ICMPTypeTimeExceeded:
		originIPPkt = icmpMsg.Body.(*icmp.TimeExceeded).Data
		// is ipv4
		if len(originIPPkt) == 0 || (originIPPkt[0]>>4) != 4 {
			return
		}

		hLen := (originIPPkt[0] & 0x0F) * 4
		srcIP := net.IPv4(originIPPkt[12], originIPPkt[13], originIPPkt[14], originIPPkt[15])
		dstIP := net.IPv4(originIPPkt[16], originIPPkt[17], originIPPkt[18], originIPPkt[19])
		// ip identifier
		ipID := binary.BigEndian.Uint16(originIPPkt[4:6])
		srcPort := binary.BigEndian.Uint16(originIPPkt[hLen : hLen+2])
		dstPort := binary.BigEndian.Uint16(originIPPkt[hLen+2 : hLen+4])
		if r.Task.DstIP == dstIP.To4().String() {
			pid := enviroments.PID
			golog.Debugf("srcIP: %s, dstIP: %s, pid: %d, deviceAddr: %s", srcIP.To4().String(), dstIP.To4().String(), pid, devIP)

			UpdateTCPPktRecord(&r.Stat.PktRecordBucket, dstIP.To4().String(), ipID, srcPort, dstPort, devIP, rxTs, false)
		}
	}
}

func (r *TCPSYNReceiver) doReceiveTCP() {
	golog.Debugf("UDPReceiver doReceiveTCP")

	conn, err := OpenReceiverIP4PacketConn("ip4:tcp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}

	r.TCPConn = conn

	filter := CreateBPFTCP(net.ParseIP(r.Task.DstIP), r.Task.DstPortMin, r.Task.DstPortMax, r.Task.SrcPortMin, r.Task.SrcPortMax)
	err = SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			conn.Close()
			return
		default:
			r.receiveTCPPkt(payloadBuf)
		}
	}
}

func (r *TCPSYNReceiver) receiveTCPPkt(payloadBuf []byte) {
	conn := r.TCPConn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	// ip header
	srcIP := h.Src
	dstIP := h.Dst
	ipID := uint16(h.ID)
	devIP := srcIP.To4().String()

	// tcp header
	srcPort := binary.BigEndian.Uint16(payload[0:2])
	dstPort := binary.BigEndian.Uint16(payload[2:4])
	tcpACK := binary.BigEndian.Uint32(payload[8:12])
	tcpFlags := payload[13]

	if (tcpFlags & 0x04) != 0 {
		// RST
		ipID = uint16(tcpACK - uint32(r.Task.MsgLen) - 1)
	} else if (tcpFlags & 0x012) == 0x12 {
		// SYN+ACK
		ipID = uint16(tcpACK - 1)
	} else if (tcpFlags & 0x010) == 0x10 {
		// ACK
		ipID = uint16(tcpACK - 1)
	} else {
		return
	}

	if r.Task.DstIP == srcIP.To4().String() {
		pid := enviroments.PID
		golog.Debugf("srcIP: %s, dstIP: %s, pid: %d, deviceAddr: %s, ipID: %d", srcIP.To4().String(), dstIP.To4().String(), pid, devIP, ipID)
		UpdateTCPPktRecord(&r.Stat.PktRecordBucket, devIP, ipID, dstPort, srcPort, devIP, rxTs, true)
	}

}

func UpdateTCPPktRecord(pktRecordBucket *PktRecordBucket, target string, ipID uint16, srcPort uint16, dstPort uint16, device string, rxTs int64, reachDst bool) {
	key := fmt.Sprintf("%s-%d-%d-%d", target, ipID, srcPort, dstPort)
	golog.Debugf("receiver key: %s", key)
	pktRecord := &PktRecord{
		Key:      key,
		Device:   device,
		RxTs:     rxTs,
		Received: true,
		ReachDst: reachDst,
	}

	pktRecordBucket.Update(pktRecord)
}
